from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import numpy as np
import logging

def par_wrapper(args, func=None):
    if args is None:
        args = ()
    func(*args)


def par_run(runner, all_tasks, n_par=8, desc='main', info_key=None,\
        b_thread=False, b_raise=True, b_disable_pbar=False):
    n_par = min(len(all_tasks), n_par)
    logger = logging.getLogger("example")
    logger.setLevel(logging.ERROR)
    if n_par > 1:
        # CPU密集型使用ProcessPoolExecutor，IO密集型使用ThreadPoolExecutor
        Pool = ThreadPoolExecutor if b_thread else ProcessPoolExecutor
        res = []
        with Pool(max_workers=n_par) as pool:
            futures = [pool.submit(runner, task) for task in all_tasks]
            for i, fut in enumerate(tqdm(futures, desc=f'par({n_par}) {desc}', disable=b_disable_pbar)):
                task = all_tasks[i]
                if info_key is not None:
                    if callable(info_key):
                        task = info_key(task)
                    else:
                        task = task[info_key]
                err = fut.exception()
                if err is None:
                    res.append(fut.result())
                else:
                    logger.error(f'{task} failed: {repr(err)}')
                    if b_raise:
                        raise err
                    res.append(None)
    else:
        res = []
        for task in tqdm(all_tasks, desc=desc, disable=b_disable_pbar):
            r = runner(task)
            res.append(r)
    return res

def generate_SIGNAL(num_signals, num_sensors, sensor_distance, wavelength,SNR=20, print_angle=True, len_signal=1000):

    noise_power = 10 ** (-SNR / 10)
    signal_angles = np.random.uniform(0, np.pi, num_signals)
    if print_angle:
        for i, angle in enumerate(sorted(signal_angles)):
            print(f"信号{i}的角度真值为：{np.rad2deg(angle)}")

    # 生成模拟水听器阵列信号
    array_signal = np.zeros((num_sensors, len_signal), dtype=complex)
    d_over_lambda = sensor_distance / wavelength

    for signal_idx, angle_truth in enumerate(signal_angles):
        for i in range(len_signal):
            source_signal = np.exp(1j * np.random.uniform(0, 2 * np.pi))  # 随机相位的声源信号
            noise = np.random.normal(0, np.sqrt(noise_power / 2), size=(num_sensors,)) + \
                    1j * np.random.normal(0, np.sqrt(noise_power / 2), size=(num_sensors,))
            array_signal[:, i] += np.exp(
                1j * 2 * np.pi * np.arange(num_sensors) * np.cos(
                    angle_truth) * d_over_lambda) * source_signal + noise
    return array_signal, np.sort(signal_angles)
